// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::ffi::{c_int, c_void};

use smallvec::smallvec;
use tracing::{debug, error, info};

use cudax::nccl::*;
use xgpu_common::{
    api_name::ApiFuncName,
    ipc::message::{Argument, ArgumentFlag, Request},
};

use crate::{fault_guard::virt, hook::nccl::NcclApi, thread};

pub struct NcclApiImpl;

impl NcclApi for NcclApiImpl {
    fn ncclGetVersion(&self, version: *mut c_int) -> ncclResult_t {
        info!("[Hooked] api_name: ncclGetVersion");
        let req = Request::with_args(
            ApiFuncName::FuncNcclgetversion as u64,
            smallvec![unsafe { Argument::from_mut_ptr(version, ArgumentFlag::ARG_OUT) }],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclCommDestroy(&self, comm: ncclComm_t) -> ncclResult_t {
        info!("[Hooked] api_name: ncclCommDestroy");

        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclcommdestroy as u64,
            smallvec![Argument::from_ref(
                &comm_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclCommAbort(&self, comm: ncclComm_t) -> ncclResult_t {
        info!("[Hooked] api_name: ncclCommAbort");

        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclcommabort as u64,
            smallvec![Argument::from_ref(
                &comm_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclCommSplit(
        &self,
        comm: ncclComm_t,
        color: c_int,
        key: c_int,
        newcomm: *mut ncclComm_t,
        config: *mut ncclConfig_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclCommSplit");

        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let mut newcomm_usize = unsafe { *newcomm };

        let req = Request::with_args(
            ApiFuncName::FuncNcclcommsplit as u64,
            smallvec![
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&color, ArgumentFlag::ARG_IN),
                Argument::from_ref(&key, ArgumentFlag::ARG_IN),
                Argument::from_mut(
                    &mut newcomm_usize,
                    ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT,
                ),
                unsafe { Argument::from_mut_ptr(config, ArgumentFlag::ARG_IN) },
            ],
        );
        let request_id = req.request_id();
        let res = thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed");
        unsafe {
            *newcomm = newcomm_usize as ncclComm_t;

            virt::handle_insert(*newcomm as *mut c_void, *newcomm as *mut c_void, 0)
                .expect("handle_insert failed");
            virt::req_id_vhandle_insert(request_id, *newcomm as *mut c_void)
                .expect("req_id_vhandle_insert failed");
        }
        res
    }

    fn ncclCommGetAsyncError(
        &self,
        comm: ncclComm_t,
        async_error: *mut ncclResult_t,
    ) -> ncclResult_t {
        debug!("[Hooked] api_name: ncclCommGetAsyncError");

        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclcommgetasyncerror as u64,
            smallvec![
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                unsafe { Argument::from_mut_ptr(async_error, ArgumentFlag::ARG_OUT) },
            ],
        );
        debug!("comm:{:?}", comm);
        debug!("async_err:{:p}, *{:?}", async_error, unsafe {
            *async_error
        });
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclBcast(
        &self,
        buff: *mut c_void,
        count: usize,
        datatype: ncclDataType_t,
        root: c_int,
        comm: ncclComm_t,
        stream: cudaStream_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclBcast");

        let buff_usize = virt::handle_map(buff).expect("handle_map failed") as usize;
        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;
        let req = Request::with_args(
            ApiFuncName::FuncNcclbcast as u64,
            smallvec![
                Argument::from_ref(&buff_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&count, ArgumentFlag::ARG_IN),
                Argument::from_ref(&datatype, ArgumentFlag::ARG_IN),
                Argument::from_ref(&root, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclAllReduce(
        &self,
        sendbuff: *const c_void,
        recvbuff: *mut c_void,
        count: usize,
        datatype: ncclDataType_t,
        op: ncclRedOp_t,
        comm: ncclComm_t,
        stream: cudaStream_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclAllReduce");
        debug!("sendbuff:{:p}", sendbuff);
        debug!("recvbuff:{:p}", recvbuff);

        let sendbuff_usize =
            virt::handle_map(sendbuff as *mut c_void).expect("handle_map failed") as usize;
        let recvbuff_usize = virt::handle_map(recvbuff).expect("handle_map failed") as usize;
        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclallreduce as u64,
            smallvec![
                Argument::from_ref(
                    &sendbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(
                    &recvbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(&count, ArgumentFlag::ARG_IN),
                Argument::from_ref(&datatype, ArgumentFlag::ARG_IN),
                Argument::from_ref(&op, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclAllGather(
        &self,
        sendbuff: *const c_void,
        recvbuff: *mut c_void,
        sendcount: usize,
        datatype: ncclDataType_t,
        comm: ncclComm_t,
        stream: cudaStream_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclAllGather");

        let sendbuff_usize =
            virt::handle_map(sendbuff as *mut c_void).expect("handle_map failed") as usize;
        let recvbuff_usize = virt::handle_map(recvbuff).expect("handle_map failed") as usize;
        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclallgather as u64,
            smallvec![
                Argument::from_ref(
                    &sendbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(
                    &recvbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(&sendcount, ArgumentFlag::ARG_IN),
                Argument::from_ref(&datatype, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclSend(
        &self,
        sendbuff: *const c_void,
        count: usize,
        datatype: ncclDataType_t,
        peer: c_int,
        comm: ncclComm_t,
        stream: cudaStream_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclSend");

        let sendbuff_usize =
            virt::handle_map(sendbuff as *mut c_void).expect("handle_map failed") as usize;
        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclsend as u64,
            smallvec![
                Argument::from_ref(
                    &sendbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(&count, ArgumentFlag::ARG_IN),
                Argument::from_ref(&datatype, ArgumentFlag::ARG_IN),
                Argument::from_ref(&peer, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclRecv(
        &self,
        recvbuff: *mut c_void,
        count: usize,
        datatype: ncclDataType_t,
        peer: c_int,
        comm: ncclComm_t,
        stream: cudaStream_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclRecv");

        let recvbuff_usize = virt::handle_map(recvbuff).expect("handle_map failed") as usize;
        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;
        let req = Request::with_args(
            ApiFuncName::FuncNcclrecv as u64,
            smallvec![
                Argument::from_ref(
                    &recvbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(&count, ArgumentFlag::ARG_IN),
                Argument::from_ref(&datatype, ArgumentFlag::ARG_IN),
                Argument::from_ref(&peer, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclGroupStart(&self) -> ncclResult_t {
        info!("[Hooked] api_name: ncclGroupStart");
        let req = Request::empty(ApiFuncName::FuncNcclgroupstart as u64);
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclGroupEnd(&self) -> ncclResult_t {
        debug!("[Hooked] api_name: ncclGroupEnd");
        let req = Request::empty(ApiFuncName::FuncNcclgroupend as u64);
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclGetUniqueId(&self, unique_id: *mut ncclUniqueId) -> ncclResult_t {
        info!("[Hooked] api_name: ncclGetUniqueId");
        debug!("unique_id: {:p}, *{:?}", unique_id, unsafe { *unique_id });
        let req = Request::with_args(
            ApiFuncName::FuncNcclgetuniqueid as u64,
            smallvec![unsafe { Argument::from_mut_ptr(unique_id, ArgumentFlag::ARG_OUT) }],
        );
        let res = thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed");
        debug!("unique_id: {:p}, *{:?}", unique_id, unsafe { *unique_id });
        res
    }

    fn ncclCommInitRankConfig(
        &self,
        comm: *mut ncclComm_t,
        nranks: c_int,
        comm_id: ncclUniqueId,
        rank: c_int,
        config: *mut ncclConfig_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclCommInitRankConfig");
        debug!("comm: {:p}, *{:?}", comm, unsafe { *comm });
        debug!("comm_id: {:?}", comm_id);

        let mut comm_usize = unsafe { *comm };

        let req = Request::with_args(
            ApiFuncName::FuncNcclcomminitrankconfig as u64,
            smallvec![
                Argument::from_mut(
                    &mut comm_usize,
                    ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&nranks, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_id, ArgumentFlag::ARG_IN),
                Argument::from_ref(&rank, ArgumentFlag::ARG_IN),
                unsafe { Argument::from_mut_ptr(config, ArgumentFlag::ARG_IN) },
                Argument::empty(), // NULL for waterlevel??
            ],
        );
        let request_id = req.request_id();
        let res = thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed");
        debug!("after comm: {:p}, *{:?}", comm, unsafe { *comm });

        unsafe {
            *comm = comm_usize as ncclComm_t;

            virt::handle_insert(*comm as *mut c_void, *comm as *mut c_void, 0)
                .expect("handle_insert failed");
            virt::req_id_vhandle_insert(request_id, *comm as *mut c_void)
                .expect("req_id_vhandle_insert failed");
        }

        // Todo: dump bootstrap metadata for recovery
        /* if NCCL_INIT_CNT == 0 {
            GLOBAL_RANK = rank;
        }
        let mut vh = VHandle::new(0,0,0,0);
        unsafe {
            vh.virt_handle = *comm as *mut c_void;
            NCCL_COMM_HANDLES[NCCL_INIT_CNT as usize] = vh;
        }
        //NCCL_BOOTSTRAPS.as_mut_ptr().add(cnt);
        let bs_ptr = unsafe {
            // NCCL_BOOTSTRAPS.as_mut_ptr().add(NCCL_INIT_CNT as usize) as *mut *mut BootstrapState
            &mut NCCL_BOOTSTRAPS[NCCL_INIT_CNT as usize] as *mut *mut BootstrapState
        };
        let bs_res = nccl_bootstrap_recv_state(bs_ptr, usock_get_fd() as i32);
        if bs_res == ncclResult_t_ncclSuccess as u32{
            debug!("Bootstrap handles received for nccl_init_cnt {:?}", unsafe {
                NCCL_INIT_CNT
            });
          } else {
            panic!("ncclBootstrapRecvState returned {:?}", bs_res);
          }
        NCCL_INIT_CNT += 1; */

        res
    }

    fn ncclCommInitRankConfigRecovery(
        &self,
        _comm: *mut ncclComm_t,
        _nranks: c_int,
        _comm_id: ncclUniqueId,
        _rank: c_int,
        _config: *mut ncclConfig_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclCommInitRankConfigRecovery");
        error!("[!!] only for recovery at server-side, client program never run into here!");
        panic!("program should NOT call ncclCommInitRankConfigRecovery");
    }

    fn ncclReduce(
        &self,
        sendbuff: *const c_void,
        recvbuff: *mut c_void,
        count: usize,
        datatype: ncclDataType_t,
        op: ncclRedOp_t,
        root: c_int,
        comm: ncclComm_t,
        stream: cudaStream_t,
    ) -> ncclResult_t {
        info!("[Hooked] api_name: ncclReduce");
        debug!("sendbuff:{:p}", sendbuff);
        debug!("recvbuff:{:p}", recvbuff);
        let stream_usize = stream;
        debug!("stream: {}, usize: 0x-{:x}", stream, stream_usize);

        let sendbuff_usize = virt::handle_map(sendbuff as *mut c_void).expect("handle_map failed")
            as *const c_void as usize;
        let recvbuff_usize = virt::handle_map(recvbuff).expect("handle_map failed") as usize;
        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;
        let req = Request::with_args(
            ApiFuncName::FuncNcclreduce as u64,
            smallvec![
                Argument::from_ref(
                    &sendbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(
                    &recvbuff_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT
                ),
                Argument::from_ref(&count, ArgumentFlag::ARG_IN),
                Argument::from_ref(&datatype, ArgumentFlag::ARG_IN),
                Argument::from_ref(&op, ArgumentFlag::ARG_IN),
                Argument::from_ref(&root, ArgumentFlag::ARG_IN),
                Argument::from_ref(&comm_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }

    fn ncclCommFinalize(&self, comm: ncclComm_t) -> ncclResult_t {
        info!("[Hooked] api_name: ncclCommFinalize");

        let comm_usize = virt::handle_map(comm as *mut c_void).expect("handle_map failed")
            as ncclComm_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncNcclcommfinalize as u64,
            smallvec![Argument::from_ref(
                &comm_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<ncclResult_t>(req)
            .expect("call invoke_api failed")
    }
}
